from rdkit import Chem
from rdkit import rdBase
from rdkit import RDConfig
import os

from rdkit.Chem import rdMolDescriptors as rdMD
from rdkit.Chem import AllChem

haveDescrs3D = hasattr(rdMD,'CalcAUTOCORR3D')

import time,unittest

def _gen3D(m,is3d,calculator):
    if not is3d:
        m = Chem.AddHs(m)
        ps = AllChem.ETKDG()
        ps.randomSeed = 0xf00d
        AllChem.EmbedMolecule(m,ps)
    return calculator(m)


class TestCase(unittest.TestCase):

  def setUp(self):
    self.dataDir = os.path.join(RDConfig.RDBaseDir,'Code','GraphMol',
    'Descriptors','test_data')
    self.suppl = Chem.SDMolSupplier(os.path.join(self.dataDir,'PBF_egfr.sdf'),removeHs=False)

  @unittest.skipIf(not haveDescrs3D,"3d descriptors not present")    
  def test1AUTOCORR2D(self):
      # not really a 3D descriptor, but this was added at the same time
      with open(os.path.join(self.dataDir,'auto2D.out')) as refFile:
          for i,m in enumerate(self.suppl):
              if i>10: break
              nm = m.GetProp('_Name')
              inl = refFile.readline()
              split = inl.split('\t')
              self.assertEqual(split[0],nm)
              split.pop(0)
              vs = rdMD.CalcAUTOCORR2D(m)
              for rv,nv in zip(split,vs):
                  self.assertAlmostEqual(float(rv),nv,delta=0.05)

  @unittest.skipIf(not haveDescrs3D,"3d descriptors not present")
  def test2AUTOCORR3D(self):
      with open(os.path.join(self.dataDir,'auto3D_dragon.out')) as refFile:
          for i,m in enumerate(self.suppl):
              if i>10: break
              nm = m.GetProp('_Name')
              inl = refFile.readline()
              split = inl.split('\t')
              self.assertEqual(split[0],nm)
              split.pop(0)
              vs = _gen3D(m,True,rdMD.CalcAUTOCORR3D)
              for rv,nv in zip(split,vs):
                  self.assertAlmostEqual(float(rv),nv,delta=0.05)

  @unittest.skipIf(not haveDescrs3D,"3d descriptors not present")
  def test3GETAWAY(self):
      with open(os.path.join(self.dataDir,'GETAWAY.new.out')) as refFile:
          for i,m in enumerate(self.suppl):
              if i>10: break
              nm = m.GetProp('_Name')
              inl = refFile.readline()
              split = inl.split('\t')
              self.assertEqual(split[0],nm)
              split.pop(0)
              vs = _gen3D(m,True,rdMD.CalcGETAWAY)
              for rv,nv in zip(split,vs):
                  self.assertAlmostEqual(float(rv),nv,delta=0.05)

  @unittest.skipIf(not haveDescrs3D,"3d descriptors not present")
  def test4MORSE(self):
      with open(os.path.join(self.dataDir,'MORSE.out')) as refFile:
          for i,m in enumerate(self.suppl):
              if i>10: break
              nm = m.GetProp('_Name')
              inl = refFile.readline()
              split = inl.split('\t')
              self.assertEqual(split[0],nm)
              split.pop(0)
              vs = _gen3D(m,True,rdMD.CalcMORSE)
              for rv,nv in zip(split,vs):
                  ref = float(rv)
                  self.assertTrue(ref < 1 or abs(ref - nv) / ref < 0.02)

  @unittest.skipIf(not haveDescrs3D,"3d descriptors not present")
  def test5RDF(self):
      with open(os.path.join(self.dataDir,'RDF.out')) as refFile:
          for i,m in enumerate(self.suppl):
              if i>10: break
              nm = m.GetProp('_Name')
              inl = refFile.readline()
              split = inl.split('\t')
              self.assertEqual(split[0],nm)
              split.pop(0)
              vs = _gen3D(m,True,rdMD.CalcRDF)
              for rv,nv in zip(split,vs):
                  ref = float(rv)
                  self.assertTrue(ref < 0.5 or abs(ref - nv) / ref < 0.02)

  @unittest.skipIf(not haveDescrs3D,"3d descriptors not present")
  def test6WHIM(self):
      with open(os.path.join(self.dataDir,'whim.new.out')) as refFile:
          for i,m in enumerate(self.suppl):
              if i>10: break
              nm = m.GetProp('_Name')
              inl = refFile.readline()
              split = inl.split('\t')
              self.assertEqual(split[0],nm)
              split.pop(0)
              vs = _gen3D(m,True,lambda x:rdMD.CalcWHIM(x,thresh=0.01))
              for rv,nv in zip(split,vs):
                  self.assertAlmostEqual(float(rv),nv,delta=0.01)


if(__name__=='__main__'):
  unittest.main()
